import pywt
import numpy as np
import matplotlib.pyplot as plt

def wavelet_denoise(signal, wavelet='db4', levels=4, threshold_mode='soft'):
    """
    小波高频降噪：在时频域中雕刻信号的纯净形态，保留低频信息。
    
    这个函数体现了三个核心模式：
    1. 分解：将复杂信号拆解为简单组件
    2. 筛选：在每个尺度上识别并保留重要信息
    3. 重构：将纯净的组件重新组合

    参数:
        signal (np.ndarray): 输入信号。
        wavelet (str): 使用的小波基函数名称，例如 'db4'。
        levels (int): 小波分解的层数。
        threshold_mode (str): 阈值处理模式，'soft' 或 'hard'。
    返回:
        tuple: 包含降噪后的信号、原始小波系数、阈值化后的小波系数和使用的阈值。
    """
    
    # 第一乐章：多尺度分解
    coeffs = pywt.wavedec(signal, wavelet, level=levels)
    
    # 第二乐章：智能阈值估计（使用贝叶斯风险准则）
    sigma = np.median(np.abs(coeffs[-1])) / 0.6745  # 噪声标准差估计
    threshold = sigma * np.sqrt(2 * np.log(len(signal)))  # 通用阈值
    threshold = 0.5

    # 第三乐章：分层净化 (高频去噪，保留低频)
    coeffs_thresh = list(coeffs)
    coeffs_thresh[0] = coeffs[0]  # 保留低频近似（主旋律）
    
    for i in range(1, len(coeffs)):
        if threshold_mode == 'soft':
            coeffs_thresh[i] = pywt.threshold(coeffs[i], threshold, mode='soft')
        else:
            coeffs_thresh[i] = pywt.threshold(coeffs[i], threshold, mode='hard')
    
    # 第四乐章：完美重构
    denoised_signal = pywt.waverec(coeffs_thresh, wavelet)
    
    return denoised_signal, coeffs, coeffs_thresh, threshold

def wavelet_denoise_low_frequency(signal, wavelet='db4', levels=4, threshold_mode='soft'):
    """
    小波低频去噪：对信号的近似系数（低频部分）进行阈值处理，保留高频细节。

    参数:
        signal (np.ndarray): 输入信号。
        wavelet (str): 使用的小波基函数名称，例如 'db4'。
        levels (int): 小波分解的层数。
        threshold_mode (str): 阈值处理模式，'soft' 或 'hard'。
    返回:
        tuple: 包含降噪后的信号、原始小波系数、阈值化后的小波系数和使用的阈值。
    """
    # 第一乐章：多尺度分解
    coeffs = pywt.wavedec(signal, wavelet, level=levels)

    
    # 第二乐章：智能阈值估计（使用贝叶斯风险准则）
    sigma = np.median(np.abs(coeffs[-1])) / 0.6745  # 噪声标准差估计
    threshold = sigma * np.sqrt(2 * np.log(len(signal)))  # 通用阈值
    threshold = 0.5

    # 第三乐章：分层净化 (低频去噪，保留高频)
    coeffs_thresh = list(coeffs)
    coeffs_thresh[0] = pywt.threshold(coeffs[0], threshold, mode=threshold_mode) # 对近似系数进行阈值处理
    
    for i in range(1, len(coeffs)):
        coeffs_thresh[i] = coeffs[i] # 保留高频细节
    
    # 第四乐章：完美重构
    denoised_signal = pywt.waverec(coeffs_thresh, wavelet)
    
    return denoised_signal, coeffs, coeffs_thresh, threshold
    


def _create_signal():
    t = np.linspace(0, 1, 500)
    linear_trend = 2 * t
    gaussian_peak = np.exp(-((t - 0.5) / 0.02) ** 2)

    clean_signal = linear_trend
    noisy_signal = linear_trend + gaussian_peak
    return t, clean_signal, noisy_signal

def _plot_denoising_results(t, clean_signal, noisy_signal, denoised, orig_coeffs, thresh_coeffs, threshold):
    """
    可视化降噪结果，包括时域对比、小波系数对比和频域分析。

    参数:
        t (np.ndarray): 时间序列。
        clean_signal (np.ndarray): 原始纯净信号。
        noisy_signal (np.ndarray): 含噪信号。
        denoised (np.ndarray): 降噪后的信号。
        orig_coeffs (list): 原始小波系数。
        thresh_coeffs (list): 阈值化后的小波系数。
        threshold (float): 使用的阈值。
    """
    fig, axes = plt.subplots(3, 2, figsize=(15, 12))

    # 时域对比
    axes[0, 0].plot(t, clean_signal, 'g-', label='原始信号', linewidth=2)
    axes[0, 0].plot(t, noisy_signal, 'r-', alpha=0.7, label='含噪信号')
    axes[0, 0].set_title('信号对比：纯净 vs 污染')
    axes[0, 0].legend()
    axes[0, 0].grid(True)

    axes[0, 1].plot(t, clean_signal, 'g-', label='原始信号', linewidth=2)
    axes[0, 1].plot(t, denoised, 'b--', label='降噪信号', linewidth=2)
    axes[0, 1].set_title('降噪效果：重获新生')
    axes[0, 1].legend()
    axes[0, 1].grid(True)

    # 小波系数对比
    for i, (orig, thresh) in enumerate(zip(orig_coeffs[1:3], thresh_coeffs[1:3])):
        axes[1, i].stem(orig, linefmt='r-', markerfmt='ro', basefmt=' ', label='原始系数')
        axes[1, i].stem(thresh, linefmt='b-', markerfmt='bo', basefmt=' ', label='阈值化后')
        axes[1, i].axhline(y=threshold, color='k', linestyle='--', label=f'阈值={threshold:.3f}')
        axes[1, i].axhline(y=-threshold, color='k', linestyle='--')
        axes[1, i].set_title(f'第{i+1}层细节系数')
        axes[1, i].legend()
        axes[1, i].grid(True)

    # 频谱分析
    freqs = np.fft.fftfreq(len(t), t[1] - t[0])
    fft_clean = np.fft.fft(clean_signal)
    fft_noisy = np.fft.fft(noisy_signal)
    fft_denoised = np.fft.fft(denoised)

    axes[2, 0].semilogy(freqs[:len(freqs)//2], np.abs(fft_noisy[:len(freqs)//2]), 'r-', alpha=0.7, label='含噪信号')
    axes[2, 0].semilogy(freqs[:len(freqs)//2], np.abs(fft_clean[:len(freqs)//2]), 'g-', linewidth=2, label='原始信号')
    axes[2, 0].set_title('频域：噪声的指纹')
    axes[2, 0].set_xlabel('频率 (Hz)')
    axes[2, 0].set_ylabel('幅度')
    axes[2, 0].legend()
    axes[2, 0].grid(True)

    axes[2, 1].semilogy(freqs[:len(freqs)//2], np.abs(fft_denoised[:len(freqs)//2]), 'b-', linewidth=2, label='降噪信号')
    axes[2, 1].semilogy(freqs[:len(freqs)//2], np.abs(fft_clean[:len(freqs)//2]), 'g--', linewidth=2, label='原始信号')
    axes[2, 1].set_title('频域：重获纯净')
    axes[2, 1].set_xlabel('频率 (Hz)')
    axes[2, 1].set_ylabel('幅度')
    axes[2, 1].legend()
    axes[2, 1].grid(True)

    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 更好的中文字体，可以显示“-”号
    plt.tight_layout()
    #plt.savefig('wavelet_denoising_analysis.png', dpi=300, bbox_inches='tight')
    plt.show()

# 演示：创造一个含噪信号并净化它
def demonstrate_denoising():
    # 创建测试信号：正弦波 + 噪声
    t, clean_signal, noisy_signal = _create_signal()

  # 应用小波降噪
    # 示例：高频去噪（保留低频）
    denoised_high_freq_removed, orig_coeffs_hf, thresh_coeffs_hf, threshold_hf = wavelet_denoise(
        noisy_signal, wavelet='haar', levels=2, threshold_mode='hard'
    )
    print("\n--- 高频去噪（保留低频）结果 ---")
    _plot_denoising_results(t, clean_signal, noisy_signal, denoised_high_freq_removed, orig_coeffs_hf, thresh_coeffs_hf, threshold_hf)
    mse_noisy_hf = np.mean((noisy_signal - clean_signal)**2)
    mse_denoised_hf = np.mean((denoised_high_freq_removed - clean_signal)**2)
    snr_improvement_hf = 10 * np.log10(mse_noisy_hf / mse_denoised_hf)
    print(f"   原始信噪比提升: {snr_improvement_hf:.2f} dB")
    print(f"   均方误差减少: {(1 - mse_denoised_hf/mse_noisy_hf)*100:.1f}%")
    print(f"   使用阈值: {threshold_hf:.4f}")

    # 示例：低频去噪（保留高频）
    denoised_low_freq_removed, orig_coeffs_lf, thresh_coeffs_lf, threshold_lf = wavelet_denoise_low_frequency(
        noisy_signal, wavelet='haar', levels=2, threshold_mode='hard'
    )
    print("\n--- 低频去噪（保留高频）结果 ---")
    _plot_denoising_results(t, clean_signal, noisy_signal, denoised_low_freq_removed, orig_coeffs_lf, thresh_coeffs_lf, threshold_lf)
    mse_noisy_lf = np.mean((noisy_signal - clean_signal)**2)
    mse_denoised_lf = np.mean((denoised_low_freq_removed - clean_signal)**2)
    snr_improvement_lf = 10 * np.log10(mse_noisy_lf / mse_denoised_lf)
    print(f"   原始信噪比提升: {snr_improvement_lf:.2f} dB")
    print(f"   均方误差减少: {(1 - mse_denoised_lf/mse_noisy_lf)*100:.1f}%")
    print(f"   使用阈值: {threshold_lf:.4f}")

    return denoised_high_freq_removed, clean_signal, noisy_signal

if __name__ == "__main__":
    demonstrate_denoising()